# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-07-08 6:47
# File     : performance_tests.py
# Project  : risk-contagion-analysis
# Desc     :

import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from src.rust_bridge import SupplyChainGraph, RiskContagionSimulator, NetworkAnalyzer


def generate_test_data(n_companies, edge_probability=0.05):
    """生成测试数据"""
    companies_data = {
        'id': [f'C{i:05d}' for i in range(1, n_companies + 1)],
        'name': [f'Company {i}' for i in range(1, n_companies + 1)],
        'risk_score': np.random.uniform(0, 1, n_companies),
        'industry': np.random.choice(['Tech', 'Manufacturing', 'Retail', 'Finance'], n_companies),
        'size': np.random.randint(50, 5000, n_companies)
    }

    # 创建关系数据 - 使用稀疏连接以模拟真实网络
    relationships = []
    for i in range(1, n_companies + 1):
        for j in range(1, n_companies + 1):
            if i != j and np.random.random() < edge_probability:
                relationships.append({
                    'source_id': f'C{i:05d}',
                    'target_id': f'C{j:05d}',
                    'strength': np.random.uniform(0.1, 0.9),
                    'relationship_type': np.random.choice(['supplier', 'customer', 'partner'])
                })

    return pd.DataFrame(companies_data), pd.DataFrame(relationships)


def run_performance_test(sizes, edge_probability=0.05):
    """针对不同网络规模运行性能测试"""

    build_times = []
    simulation_times = []
    analysis_times = []

    for size in sizes:
        print(f"测试网络规模: {size} 公司")

        # 生成数据
        companies_df, relationships_df = generate_test_data(size, edge_probability)

        # 测试网络构建时间
        start_time = time.time()
        graph = SupplyChainGraph()
        graph.build_from_dataframe(companies_df, relationships_df)
        build_time = time.time() - start_time
        build_times.append(build_time)
        print(f"  构建时间: {build_time:.4f} 秒")

        # 测试模拟时间
        simulator = RiskContagionSimulator()
        start_time = time.time()
        simulator.simulate(graph)
        sim_time = time.time() - start_time
        simulation_times.append(sim_time)
        print(f"  模拟时间: {sim_time:.4f} 秒")

        # 测试分析时间
        analyzer = NetworkAnalyzer()
        start_time = time.time()
        analyzer.calculate_centrality(graph)
        analysis_time = time.time() - start_time
        analysis_times.append(analysis_time)
        print(f"  分析时间: {analysis_time:.4f} 秒")

    # 绘制性能图表
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.plot(sizes, build_times, 'o-')
    plt.title('网络构建时间')
    plt.xlabel('公司数量')
    plt.ylabel('时间 (秒)')

    plt.subplot(1, 3, 2)
    plt.plot(sizes, simulation_times, 'o-')
    plt.title('风险模拟时间')
    plt.xlabel('公司数量')
    plt.ylabel('时间 (秒)')

    plt.subplot(1, 3, 3)
    plt.plot(sizes, analysis_times, 'o-')
    plt.title('网络分析时间')
    plt.xlabel('公司数量')
    plt.ylabel('时间 (秒)')

    plt.tight_layout()
    plt.savefig("performance_results.png")
    plt.show()

    return {
        'sizes': sizes,
        'build_times': build_times,
        'simulation_times': simulation_times,
        'analysis_times': analysis_times
    }


if __name__ == "__main__":
    # 测试不同规模的网络
    network_sizes = [10, 50, 100, 500, 1000]
    results = run_performance_test(network_sizes, edge_probability=0.01)

    # 打印结果表格
    print("\n性能测试结果汇总:")
    print("-" * 80)
    print(f"{'网络规模':<12}{'构建时间 (秒)':<20}{'模拟时间 (秒)':<20}{'分析时间 (秒)':<20}")
    print("-" * 80)

    for i, size in enumerate(results['sizes']):
        print(
            f"{size:<12}{results['build_times'][i]:<20.4f}{results['simulation_times'][i]:<20.4f}{results['analysis_times'][i]:<20.4f}")
